import os
import sys
import onnx
import onnx_graphsurgeon as gs

# === 工程路径配置 ===
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)

# ======= 保存文件夹设置 =======
folder = 'split_merge_onnx'
os.makedirs(folder, exist_ok=True)

# 路径配置
BACKBONE_ONNX =  os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_int8_ALLConcat.onnx')
NMS_ONNX      =  os.path.join(PROJECT_ROOT, 'onnx/yolov11n_nms.onnx')
MERGED_ONNX   =  os.path.join(PROJECT_ROOT, 'engine/quantization_onnx/yolov11n_nms_merge.onnx')

MERGE_INSERT_NODE = "/Transpose"

print("载入主干(on), NMS(on)...")
backbone_graph = gs.import_onnx(onnx.load(BACKBONE_ONNX))
nms_graph = gs.import_onnx(onnx.load(NMS_ONNX))

# 1. 获取 BACKBONE 原输出名
if len(backbone_graph.outputs) != 1:
    raise RuntimeError("主干输出不只一个，请手动指定")
BACKBONE_OUTPUT_NAME = backbone_graph.outputs[0].name  # 例如 'output0'
print("主干输出名:", BACKBONE_OUTPUT_NAME)
NMS_OUTPUT_NEW_NAME = 'nms_output0'  # 合并后新输出名

# 2. 移除 backbone 的输出 output0
print(f"移除主干模型输出 {BACKBONE_OUTPUT_NAME}")
backbone_graph.outputs.clear()

# 3. 找到 NMS ONNX 里的“/Transpose”节点，以及其后全部节点
def collect_downstream_nodes(start_node, nodes):
    """从某node起，BFS遍历收集所有下游节点"""
    selected = set()
    queue = [start_node]
    while queue:
        node = queue.pop()
        node_id = id(node)
        if node_id in selected:
            continue
        selected.add(node_id)
        for out in node.outputs:
            for n in nodes:
                if out in n.inputs:
                    queue.append(n)
    return [n for n in nodes if id(n) in selected]

# NMS里/Transpose节点（注意前缀）
transpose_candidates = [n for n in nms_graph.nodes if n.name == (MERGE_INSERT_NODE)]
assert len(transpose_candidates) == 1, "NMS模型未找到/Transpose节点"
transpose_node = transpose_candidates[0]
nms_nodes_to_merge = collect_downstream_nodes(transpose_node, nms_graph.nodes)
print(f"NMS子图拼接节点数量: {len(nms_nodes_to_merge)}")

# 4. *********************
#    NMS子图里所有output0（包括outputs数组、node.inputs/outputs）进行重命名，避免与backbone冲突
# *********************
def rename_nms_variables(nodes, nms_outputs, old_name, new_name):
    """把NMS子图内所有变量名old_name替换为new_name"""
    # 所有变量对象
    for node in nodes:
        # 重命名 outputs
        for i, v in enumerate(node.outputs):
            if v.name == old_name:
                node.outputs[i].name = new_name
        # 重命名 inputs
        for i, v in enumerate(node.inputs):
            if v.name == old_name:
                node.inputs[i].name = new_name
    # outputs对象单独重命名
    for i, v in enumerate(nms_outputs):
        if v.name == old_name:
            nms_outputs[i].name = new_name

rename_nms_variables(nms_nodes_to_merge, nms_graph.outputs, BACKBONE_OUTPUT_NAME, NMS_OUTPUT_NEW_NAME)

# 5. 替换NMS子图/Transpose节点输入为主干输出 Variable对象
nms_trans_input_name = transpose_node.inputs[0].name
backbone_output_var = None
for var in backbone_graph.tensors().values():
    if var.name == BACKBONE_OUTPUT_NAME:
        backbone_output_var = var
        break
assert backbone_output_var is not None

# 确保NMS的/Transpose输入直接对应backbone输出
for node in nms_nodes_to_merge:
    node.inputs = [backbone_output_var if i.name == nms_trans_input_name else i for i in node.inputs]

# 6. 设置合并后新输出
print(f"设置整体输出变量为: {NMS_OUTPUT_NEW_NAME}")
new_outputs = []
for v in nms_graph.outputs:
    if v.name == NMS_OUTPUT_NEW_NAME:
        new_outputs.append(v)
backbone_graph.outputs = new_outputs

# 7. 合并节点
backbone_graph.nodes.extend(nms_nodes_to_merge)

# 8. 合并 NMS 常量（只合并未存在的常量）
backbone_tensor_names = set(backbone_graph.tensors())
for t in nms_graph.tensors().values():
    if isinstance(t, gs.Constant) and t.name not in backbone_tensor_names:
        backbone_graph.tensors()[t.name] = t

# 9. 清理和保存
print("图合并完成，清理中...")
backbone_graph.cleanup()
onnx.save(gs.export_onnx(backbone_graph), MERGED_ONNX)
print(f"合并模型已保存到 {MERGED_ONNX}")